import numpy as np

from data_preparation.trajectory_set import TrajectorySet
from tools.data_reader import DataReader
from tools.data_writer import DataWriter
import config.folder_and_file_names as config
from generator.trajectory_generator import Generator
from generator.state_trajectory_generation import StateGeneration
from discretization.grid import Grid
# from tools.object_store import ObjectStore
from tools.general_tools import GeneralTools


class RealLocationTranslator:

    def __init__(self, cc):
        self.grid = Grid(cc)

    # this function gives translator grid to use
    def load_translator(self, grid):
        self.grid = grid

    def translate_given_state_sequence(self, state_sequence):
        grid1 = self.grid
        if len(state_sequence) < 2:
            start_end_array = np.empty(2, dtype=int)
            start_end_array[0] = state_sequence[0]
            start_end_array[1] = state_sequence[0]
            state_sequence = start_end_array
        all_level2_state_borders = grid1.level2_cell_borders
        trajectory_length = len(state_sequence)
        real_trajectory = np.random.random((trajectory_length, 2))#生成一个形状为 (trajectory_length, 2) 的二维数组
        for index_of_states in range(trajectory_length):
            state = state_sequence[index_of_states]
            borders = all_level2_state_borders[state]
            location = self.sample_from_a_subcell(borders)
            real_trajectory[index_of_states, :] = location
        return real_trajectory

    def translate_given_state_sequence_by_tree(self, state_sequence):
        grid1 = self.grid
        if len(state_sequence) < 2:
            start_end_array = np.empty(2, dtype=int)
            start_end_array[0] = state_sequence[0]
            start_end_array[1] = state_sequence[0]
            state_sequence = start_end_array
        all_level2_state_borders = grid1.level2_cell_borders
        trajectory_length = len(state_sequence)
        real_trajectory = np.random.random((trajectory_length, 2))  # 生成一个形状为 (trajectory_length, 2) 的二维数组
        for index_of_states in range(trajectory_length):
            state = state_sequence[index_of_states]
            # print("state:", state,type(state))
            borders = all_level2_state_borders[state]
            location = self.sample_from_a_subcell(borders)
            real_trajectory[index_of_states, :] = location
        return real_trajectory

    def sample_from_a_subcell(self, borders):
        gt1 = GeneralTools()
        north = borders[0]
        south = borders[1]
        west = borders[2]
        east = borders[3]
        x_value = gt1.sample_from_interval(west, east)
        y_value = gt1.sample_from_interval(south, north)
        location = np.array([x_value, y_value])
        return location

    def get_real_trajectories(self, state_trajectories,point_trajectory_set: TrajectorySet):
        trajectory_number = point_trajectory_set.get_trajectory_number()
        for i in range(trajectory_number):
            trajectory = point_trajectory_set.give_trajectory_by_index(i)
            index_array = trajectory.level2_unique_index_sequence.copy()
            experiment_trajectory_array = self.translate_given_state_sequence(index_array)
            trajectory.experiment_array = experiment_trajectory_array
        ####
        state_trajectory_list = state_trajectories
        real_trajectory_list = []
        for state_trajectory in state_trajectory_list:
            real_trajectory = self.translate_given_state_sequence(state_trajectory)
            real_trajectory_list.append(real_trajectory)
        return real_trajectory_list

    def translate_trajectories(self, grid, state_trajectories,point_trajectory_set: TrajectorySet):
        self.load_translator(grid)
        real_tra = self.get_real_trajectories(state_trajectories,point_trajectory_set)
        return real_tra

    def translate_trajectories_by_tree(self, grid, state_trajectories):
        self.load_translator(grid)
        state_trajectory_list = state_trajectories
        real_trajectory_list = []
        for state_trajectory in state_trajectory_list:
            real_trajectory = self.translate_given_state_sequence_by_tree(state_trajectory)
            real_trajectory_list.append(real_trajectory)
        return real_trajectory_list
